"""
Script to post-process the lmp data file generate from LigParGen web server (http://zarbi.chem.yale.edu/ligpargen/)
and prepare a molecule file for Alkene.
"""
def Process(lpgfile, molfile, neutralize):
    """
    To post process the LigParGen lmp data file:
    1) Atom types were reassigned by treating atoms with same mass and pair-coeffs as same type;
    2) Bond types with same bond parameters are seen as same type;
    3) Angle types with same angle parameters are seen as same type;
    4) Dihedral types with same dihedral parameter are seen as same type;
    5) Impropers were deleted.
    """
    import pandas as pd
    from pymatgen.io.lammps.data import LammpsData
    from pymatgen.core.periodic_table import Element, get_el_sp

    data = LammpsData.from_file(lpgfile)

    # Delete impropers
    if 'Impropers' in data.topology.keys():
       data.topology.pop("Impropers")
    if 'Improper Coeffs' in data.force_field.keys():
       data.force_field.pop("Improper Coeffs")

    # zero net charge
    if neutralize:
       data.atoms.q = data.atoms.q - data.atoms.q.sum()/len(data.atoms.q)

    # Identify atom types
    mass_and_pair = data.masses.join(data.force_field['Pair Coeffs'])

    tflag = list()
    for index,row in mass_and_pair.iterrows():
        tflag.append((row.mass, row.coeff1, row.coeff2))
    tflag_set = set(tflag)

    flag2type = dict()
    for index, flag in enumerate(tflag_set, start=1):
        flag2type[flag] = index

    print("The identified atom types are:\n Type   Mass   PairCoeff")
    for flag, index in flag2type.items():
        print(f"   {index} : {flag}")
    ntype = len(flag2type)

    # User can customize the atom types
    str_in = input(f"If you would like to redefine the atom types, input the ids in sequense, Enter to accept: ")
    user_types = str_in.split()
    if len(user_types) == ntype:
       for f,i in flag2type.items():
           flag2type[f] = int(user_types[i-1])

       print("The reassigned atom types are:\n Type   Mass   PairCoeff")
       for flag, index in flag2type.items():
           print(f"   {index} : {flag}")

    masses = [None] * (ntype+1)
    coeff1 = [None] * (ntype+1)
    coeff2 = [None] * (ntype+1)

    # Reset Masses and Pair Coeff info
    for flag, index in flag2type.items():
        m,c1,c2 = flag
        masses[index] = m
        coeff1[index] = c1
        coeff2[index] = c2
    
    data.masses = pd.DataFrame({'mass':masses}).drop(0)
    data.force_field['Pair Coeffs'] = pd.DataFrame({'coeff1':coeff1, 'coeff2':coeff2}).drop(0)

    # Reset atom types for all atoms
    type_old2new = dict()
    for index, flag in enumerate(tflag, start=1):
        type_old2new[index] = flag2type[flag]

    data.atoms['type'] = data.atoms['type'].map(lambda x:type_old2new[x])

    # Redefine bond types
    bflag = list()
    for index, row in data.force_field['Bond Coeffs'].iterrows():
        bflag.append((row.coeff1, row.coeff2))
    bflag_set = set(bflag)

    bflag2type = dict()
    for index, flag in enumerate(bflag_set, start=1):
        bflag2type[flag] = index

    nbtype = len(bflag2type)
    coeff1 = [None] * (nbtype+1)
    coeff2 = [None] * (nbtype+1)

    # Reset bond coeffs
    for flag, index in bflag2type.items():
        c1,c2 = flag
        coeff1[index] = c1
        coeff2[index] = c2
    data.force_field['Bond Coeffs'] = pd.DataFrame({"coeff1":coeff1, "coeff2":coeff2}).drop(0)

    print("\nThe identified bond types are:\n Type   BondCoeff")
    for flag, index in bflag2type.items():
        print(f"   {index} : {flag}")

    # Reset bond types
    btype_old2new = dict()
    for index, flag in enumerate(bflag, start=1):
        btype_old2new[index] = bflag2type[flag]

    data.topology['Bonds']['type'] = data.topology['Bonds']['type'].map(lambda x:btype_old2new[x])

    # Redefine angle types
    aflag = list()
    for index, row in data.force_field['Angle Coeffs'].iterrows():
        aflag.append((row.coeff1, row.coeff2))
    aflag_set = set(aflag)

    aflag2type = dict()
    for index, flag in enumerate(aflag_set, start=1):
        aflag2type[flag] = index

    natype = len(aflag2type)
    coeff1 = [None] * (natype+1)
    coeff2 = [None] * (natype+1)

    # Reset angle coeffs
    for flag, index in aflag2type.items():
        c1,c2 = flag
        coeff1[index] = c1
        coeff2[index] = c2
    data.force_field['Angle Coeffs'] = pd.DataFrame({"coeff1":coeff1, "coeff2":coeff2}).drop(0)

    print("\nThe identified angle types are:\n Type   AngleCoeff")
    for flag, index in aflag2type.items():
        print(f"   {index} : {flag}")

    # Reset angle types
    atype_old2new = dict()
    for index, flag in enumerate(aflag, start=1):
        atype_old2new[index] = aflag2type[flag]

    data.topology['Angles']['type'] = data.topology['Angles']['type'].map(lambda x:atype_old2new[x])

    # Redefine Dihedral types
    dflag = list()
    for index, row in data.force_field['Dihedral Coeffs'].iterrows():
        dflag.append((row.coeff1, row.coeff2, row.coeff3, row.coeff4))
    dflag_set = set(dflag)

    dflag2type = dict()
    for index, flag in enumerate(dflag_set, start=1):
        dflag2type[flag] = index

    ndtype = len(dflag2type)
    coeff1 = [None] * (ndtype+1)
    coeff2 = [None] * (ndtype+1)
    coeff3 = [None] * (ndtype+1)
    coeff4 = [None] * (ndtype+1)

    # Reset dihedral coeffs
    for flag, index in dflag2type.items():
        c1,c2,c3,c4 = flag
        coeff1[index] = c1
        coeff2[index] = c2
        coeff3[index] = c3
        coeff4[index] = c4
    data.force_field['Dihedral Coeffs'] = pd.DataFrame({"coeff1":coeff1, "coeff2":coeff2, "coeff3":coeff3, "coeff4":coeff4}).drop(0)

    print("\nThe identified dihedral types are:\n Type   DihedralCoeff")
    for flag, index in dflag2type.items():
        print(f"   {index} : {flag}")

    # Reset dihedral types
    dtype_old2new = dict()
    for index, flag in enumerate(dflag, start=1):
        dtype_old2new[index] = dflag2type[flag]

    data.topology['Dihedrals']['type'] = data.topology['Dihedrals']['type'].map(lambda x:dtype_old2new[x])

    # Map element names from atomic mass
    ename = dict()
    for index, mass in enumerate(data.masses['mass'], start=1):
        ename[index] = 'X'
        diff_mass = mass*mass
        for elem in Element:
            e = get_el_sp(elem)
            dm = (mass - e.atomic_mass)**2
            if dm < diff_mass:
               diff_mass = dm
               ename[index] = e.symbol
        
    # output the file read by alkene
    lines = []
    lines.append(f"Molecule_from_{lpgfile}")
    lines.append(f"{data.structure.formula}".replace(" ", "_"))
    lines.append(f"{len(data.masses)} {len(data.force_field['Bond Coeffs'])} {len(data.force_field['Angle Coeffs'])} {len(data.force_field['Dihedral Coeffs'])}")
    lines.append(f"{len(data.atoms)} {len(data.topology['Bonds'])} {len(data.topology['Angles'])} {len(data.topology['Dihedrals'])}")
    lines.append("1.0 4\n\nElemental Info")
    for index,mc in enumerate(zip(data.masses['mass'], data.force_field['Pair Coeffs'].coeff1, data.force_field['Pair Coeffs'].coeff2), start=1):
        mass,c1,c2 = mc
        lines.append(f"{index} {ename[index]} {mass} {c2} {c1} 1.0")

    lines.append("\nBond Coeffs")
    for index,coeff in data.force_field['Bond Coeffs'].iterrows():
        lines.append(f"{index} {coeff.coeff1} {coeff.coeff2}")

    lines.append("\nAngle Coeffs")
    for index,coeff in data.force_field['Angle Coeffs'].iterrows():
        lines.append(f"{index} {coeff.coeff1} {coeff.coeff2}")

    lines.append("\nDihedral Coeffs")
    for index,coeff in data.force_field['Dihedral Coeffs'].iterrows():
        lines.append(f"{index} {coeff.coeff1} {coeff.coeff2} {coeff.coeff3} {coeff.coeff4}")

    lines.append("\nAtoms")
    for index,atom in data.atoms.iterrows():
        lines.append(f"{index} {int(atom['molecule-ID'])} {int(atom.type)} {atom.q} {atom.x} {atom.y} {atom.z}")

    lines.append("\nBonds")
    for index,bond in data.topology['Bonds'].iterrows():
        lines.append(f"{index} {bond.type} {bond.atom1} {bond.atom2}")

    lines.append("\nAngles")
    for index,angle in data.topology['Angles'].iterrows():
        lines.append(f"{index} {angle.type} {angle.atom1} {angle.atom2} {angle.atom3}")

    lines.append("\nDihedrals")
    for index,dihed in data.topology['Dihedrals'].iterrows():
        lines.append(f"{index} {dihed.type} {dihed.atom1} {dihed.atom2} {dihed.atom3} {dihed.atom4}")

    fp = open(molfile, "w")
    fp.write("\n".join(lines))
    fp.close()

    print(f"\nProcessed infomration on molecule {data.structure.formula} written to {molfile}.")
    return


if __name__ == "__main__":
   import os
   import argparse
   parser = argparse.ArgumentParser()
   parser.add_argument("file", help="The lmp data file from LigParGen to analyze.")
   parser.add_argument("-o", "--output", default="molecule.dat", help="To define the output file name; by default: molecule.dat.")
   parser.add_argument("-c", "--charge", help="To reset the net charge to zero. By default: not set.", action="count")

   args = parser.parse_args()
    
   if not os.path.exists(args.file):
      print(f"\nFile {args.file} required but not found!\n")

   Process(args.file, args.output, args.charge)
